from typing import Annotated, Optional

from langchain_community.tools import TavilySearchResults
from langchain_core.prompts import ChatPromptTemplate
from typing_extensions import TypedDict
from langgraph.graph.message import AnyMessage, add_messages

from langchain_openai import ChatOpenAI
from langchain_community.chat_models.tongyi import ChatTongyi

from langgraph_customer.tools.utilities_tools import _print_event
from langgraph_customer.utils.update_sqlite import update_table_data
from langchain_core.runnables import Runnable, RunnableConfig

from langgraph_customer.utils.config import Settings
from langgraph_customer.tools.flight_tools import *
from langgraph_customer.tools.car_tools import *
from langgraph_customer.tools.hotel_tools import *
from langgraph_customer.tools.utilities_tools import *
from langgraph_customer.tools.excursions_tools import *

class State(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]

class Assistant:
    def __init__(self, runnable: Runnable):
        self.runnable = runnable

    def __call__(self, state: State, config: RunnableConfig):
        while True:
            configuration = config.get("configurable", {})
            passenger_id = configuration.get("passenger_id", None)
            state = {**state, "user_info": passenger_id}
            result = self.runnable.invoke(state)
            # If the LLM happens to return an empty response, we will re-prompt it
            # for an actual response.
            if not result.tool_calls and (
                not result.content
                or isinstance(result.content, list)
                and not result.content[0].get("text")
            ):
                messages = state["messages"] + [("user", "用真实的输出做出回应。")]
                state = {**state, "messages": messages}
            else:
                break
        return {"messages": result}

llm = ChatTongyi(
    model="qwen-max",
    api_key=Settings.OPENAI_API_KEY,
)
# llm = ChatOpenAI(model=Settings.LLM_NAME, base_url=Settings.LLM_URL)

primary_assistant_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """您是瑞士航空的客服助理，乐于助人。使用提供的工具搜索航班、公司政策和其他信息，以协助用户查询。搜索时，请坚持不懈。如果第一次搜索没有结果，请扩大搜索范围。如果搜索结果为空，请在放弃之前扩大搜索范围。
    当前用户：<User>\n{user_info}\n</User>
    当前时间：{time}。
            """
        ),
        ("placeholder", "{messages}"),
    ]
).partial(time=datetime.now)


part_1_tools = [
    TavilySearchResults(max_results=1),
    fetch_user_flight_information,
    search_flights,
    lookup_policy,
    update_ticket_to_new_flight,
    cancel_ticket,
    search_car_rentals,
    book_car_rental,
    update_car_rental,
    cancel_car_rental,
    search_hotels,
    book_hotel,
    update_hotel,
    cancel_hotel,
    search_trip_recommendations,
    book_excursion,
    update_excursion,
    cancel_excursion,
]

part_1_assistant_runnable = primary_assistant_prompt | llm.bind_tools(part_1_tools)

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import END, StateGraph, START
from langgraph.prebuilt import tools_condition

builder = StateGraph(State)


# Define nodes: these do the work
builder.add_node("assistant", Assistant(part_1_assistant_runnable))
builder.add_node("tools", create_tool_node_with_fallback(part_1_tools))
# Define edges: these determine how the control flow moves
builder.add_edge(START, "assistant")
builder.add_conditional_edges(
    "assistant",
    tools_condition,
)
builder.add_edge("tools", "assistant")

# The checkpointer lets the graph persist its state
# this is a complete memory for the entire graph.
memory = InMemorySaver()
part_1_graph = builder.compile(checkpointer=memory)

from IPython.display import Image, display

try:
    display(Image(part_1_graph.get_graph(xray=True).draw_mermaid_png()))
except Exception:
    # This requires some extra dependencies and is optional
    pass

import shutil
import uuid

# Let's create an example conversation a user might have with the assistant
tutorial_questions = [
    "您好，我的航班几点？",
    "我可以把航班改早吗？我想今天晚点出发。",
    "那请把我的航班改到下周的某个时间。",
    "下一个可选方案很棒。",
    "住宿和交通怎么样？",
    "是的，我想预订一家经济实惠的酒店，作为我为期一周（7天）的住宿。我还想租一辆车。",
    "好的，您能预订一下您推荐的酒店吗？听起来不错。",
    "好的，请预订任何价格适中且有空房的酒店。",
    "现在说到租车，我有什么选择？",
    "太棒了，我们就选最便宜的吧。预订7天的。",
    "太棒了，那么您现在有什么短途旅行推荐吗？",
    "我去的时候还有空房吗？",
    "有意思——我喜欢博物馆，有什么选择吗？",
    "好的，选一个，预订我第二天去那里的。"
]


# 修改表信息符合现在情况
db = update_table_data("../resource/travel2.sqlite")
thread_id = str(uuid.uuid4())

config = {
    "configurable": {
        "passenger_id": "3442 587242",
        "thread_id": thread_id,
    }
}


_printed = set()
for question in tutorial_questions:
    events = part_1_graph.stream(
        {"messages": ("user", question)}, config, stream_mode="values"
    )
    for event in events:
        _print_event(event, _printed)








